// Copyright 2022 Chen Jun
// 在MIT许可证下发布。

#ifndef ARMOR_DETECTOR_ARMOR_HPP_
#define ARMOR_DETECTOR_ARMOR_HPP_

#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <Eigen/Eigen>

// STL
#include <algorithm>
#include <string>

namespace hnurm
{
    const int RED = 0;
    const int BLUE = 1;
    const int PURPLE = 2;


    enum ArmorType
    {
        SMALL = 0, LARGE = 1
    };

    /**
     * @brief Light类表示一个灯，继承自cv::RotatedRect类
     */
    struct Light: public cv::RotatedRect
    {
        Light() = default;

        /**
         * @brief 构造函数
         * @param box 矩形框
         */
        explicit Light(const cv::RotatedRect& box)
            : cv::RotatedRect(box)
        {
            cv::Point2f p[4];
            box.points(p);
            std::sort(p, p + 4, [](const cv::Point2f &a, const cv::Point2f &b)
                      { return a.y < b.y; });
            top = (p[0] + p[1]) / 2;
            bottom = (p[2] + p[3]) / 2;

            length = cv::norm(top - bottom);
            width = cv::norm(p[0] - p[1]);

            tilt_angle = std::atan2(std::abs(top.x - bottom.x), std::abs(top.y - bottom.y));
            tilt_angle = tilt_angle / CV_PI * 180;
        }

        int color;                    // 灯的颜色
        cv::Point2f top, bottom;       // 灯的顶点和底点
        double length;                 // 灯的长度
        double width;                  // 灯的宽度
        float tilt_angle;              // 灯的倾斜角度
    };

    /**
     * @brief Armor类表示一个装甲，包含灯、中心点、四个角点、数字ROI图像、数字、索引、相似度、置信度、分类结果、装甲类型等信息
     */
    struct Armor
    {
        // 左灯条、右灯条
        Light left_light, right_light;

        // 装甲板中心点
        cv::Point2f center;

        // 装甲板的四个角点，顺序为左上、左下、右下、右上
        std::vector<cv::Point2f> points2d;

        // 数字ROI图像
        cv::Mat number_img;

        std::string number;          // 数字
        int idx{};                   // 索引
        float similarity{};          // 相似度
        float confidence{};          // 置信度
        std::string classification_result;   // 分类结果

        ArmorType armor_type = ArmorType::SMALL;  // 装甲类型

        double distance_to_image_center{};      // 距离图像中心的像素值
        // todo
        Eigen::Vector3d position;           // 位置
        Eigen::Matrix3d rotation;            // 旋转
        Eigen::Matrix3d _r1;                 // 相机坐标系到云台坐标系
        Eigen::Matrix3d _r2;                 // 先对pitch进行旋转
        Eigen::Matrix3d _r3;                 // 再对yaw进行旋转，变为世界坐标系（pitch、yaw为0）
        Eigen::Matrix3d _rmat;
        Eigen::Vector3d _translation;

        Armor() = default;

        /**
         * @brief 构造函数
         * @param l1 左灯条
         * @param l2 右灯条
         */
        Armor(const Light &l1, const Light &l2)
        {
            if (l1.center.x < l2.center.x)
            {
                left_light = l1, right_light = l2;
            }
            else
            {
                left_light = l2, right_light = l1;
            }
            center = (left_light.center + right_light.center) / 2;

            points2d.resize(9);
            points2d[0] = left_light.top;
            points2d[1] = left_light.bottom;
            points2d[2] = right_light.bottom;
            points2d[3] = right_light.top;

            points2d[4] = (left_light.top+right_light.top)/2;
            points2d[5] = (left_light.top+left_light.bottom)/2;
            points2d[6] = (left_light.bottom+right_light.bottom)/2;
            points2d[7] = (right_light.bottom+right_light.top)/2;

            points2d[8] = center;
        }

        /**
         * @brief 在给定的图像上绘制装甲
         * @param canvas 绘制的图像
         */
        void showArmor(cv::Mat &canvas)
        {
            // 绘制装甲
            cv::line(canvas, points2d[0], points2d[1], cv::Scalar(0, 255, 0), 2);
            cv::line(canvas, points2d[1], points2d[2], cv::Scalar(0, 255, 0), 2);
            cv::line(canvas, points2d[2], points2d[3], cv::Scalar(0, 255, 0), 2);
            cv::line(canvas, points2d[3], points2d[0], cv::Scalar(0, 255, 0), 2);
            // 绘制中心点
            cv::circle(canvas, center, 10, cv::Scalar(0, 255, 0), 2);
            // 在装甲的左上角绘制分类结果
            cv::putText(canvas, classification_result, points2d[0], cv::FONT_HERSHEY_SIMPLEX,
                        1, cv::Scalar(0, 255, 0), 2);

        }
    };

    /**
     * @brief Center类表示能量机关中心R，继承自cv::RotatedRect类
     */
    struct Center: public cv::RotatedRect
    {
        Center() = default;

        /**
         * @brief 构造函数
         * @param box 矩形框
         */
        explicit Center(const cv::RotatedRect& box)
            : cv::RotatedRect(box)
        {
            cv::Point2f p[4];
            box.points(p);
            std::sort(p, p + 4, [](const cv::Point2f &a, const cv::Point2f &b)
                      { return a.y < b.y; });
            top = (p[0] + p[1]) / 2;
            bottom = (p[2] + p[3]) / 2;

            length = cv::norm(top - bottom);
            width = cv::norm(p[0] - p[1]);
        }

        int color;                    // 灯的颜色
        cv::Point2f top, bottom;       // 灯的顶点和底点
        double length;                 // 灯的长度
        double width;                  // 灯的宽度
    };

    /**
     * @brief Flabellum类表示能量机关的扇叶，继承自cv::RotatedRect类
     */
    struct Flabellum: public cv::Rect
    {
        Flabellum() = default;

        /**
         * @brief 构造函数
         * @param box 矩形框
         */
        explicit Flabellum(const cv::Rect& box)
            : cv::Rect(box)
        {
            top = box.tl() + cv::Point_(box.width / 2, 0);
            bottom = box.br() - cv::Point_(box.width / 2, 0);

            length = box.height;
            width = box.width;

            center = (box.br() + box.tl()) / 2;
            current_state = 0;
        }

        float operator^(const Flabellum &b) {
            int x_min = std::max(this->tl().x, b.tl().x);
            int y_min = std::max(this->tl().y, b.tl().y);
            int x_max = std::min(this->br().x, b.br().x);
            int y_max = std::min(this->br().y, b.br().y);

            int iw = std::max(x_max - x_min, 0);
            int ih = std::max(y_max - y_min, 0);

            int inters = iw * ih;

            int uni = this->area() + b.area() - inters;
            // 计算两个矩形的iou
            float iou = inters * 1.0 / (uni + 1e-6);
            // 计算两个矩形的的diou
            float diou = iou - distance(center, b.center) / distance((*this|b).br(),(*this|b).tl());

            return iou;
        }

        float distance(cv::Point2f a, cv::Point2f b) {
            return (a.x - b.x)*(a.x - b.x) + (a.y - b.y)*(a.y - b.y);
        }

        cv::RotatedRect rotate;
        int color;                    // 灯的颜色
        cv::Point2f top, bottom, center;       // 灯的顶点和底点
        double length;                 // 灯的长度
        double width;                  // 灯的宽度
        int id;                        // 扇叶id
        double angle;
        int current_state;              // 当前状态0不亮1待击打2已激活

        std::vector<cv::Point2f> points2d;
        Eigen::Vector3d position;           // 位置
        Eigen::Matrix3d rotation;            // 旋转
        Eigen::Matrix3d _r1;                 // 相机坐标系到云台坐标系
        Eigen::Matrix3d _r2;                 // 先对pitch进行旋转
        Eigen::Matrix3d _r3;                 // 再对yaw进行旋转，变为世界坐标系（pitch、yaw为0）
        Eigen::Matrix3d _rmat;
        Eigen::Vector3d _translation;
    };

    /**
     * @brief Energybuff类表示能量机关
    */
    struct Energybuff{
        Energybuff() {
            // flabellums.resize(5);
            // candflabs.resize(5);
        }

        std::vector<Flabellum>flabellums;
        std::vector<Flabellum>candflabs;  //候选扇叶

        int lighted_num;
        Center center;              //中心R
        double radius;
        std::vector<Eigen::Vector3d>positions;
        Flabellum targetarmor;
        Flabellum lasttarget;
        float speed;

        //todo
        enum Type{
            small,
            large
        }type;
        enum class TrackState { LOST = 0, TRACK = 1};
        TrackState trackstate;
        
        //防止误识别 待数据稳定才更新
        int new_num;
        int update_time;


        /**
         * @brief 计算候选框的位置
        */
        void init(Flabellum flabellum) {
            cv::Point_ r = center.center;
            cv::Point_ f = flabellum.center;
            radius = sqrt((r.x - f.x)*(r.x - f.x) + (r.y - f.y)*(r.y - f.y));
            flabellum.angle = atan2(r.y-f.y, r.x-f.x);
            double angle = flabellum.angle;
            candflabs.clear();
            for(int i=0;i<5;i++) {
                int height = flabellum.height;
                int width = flabellum.width;
                double add_angle = 2 * CV_PI / 5 * i + CV_PI;
                int x = r.x + radius * cos(angle + add_angle) - width / 2;
                int y = r.y + radius * sin(angle + add_angle) - height / 2;
                
                cv::Rect rect(x, y, width, height);
                Flabellum newflab(rect);
                newflab.id = i;
                newflab.current_state = (i==0)?1:0;
                //保留旋转矩形
                if(i == 0) {
                    newflab.rotate = flabellum.rotate;
                }
                candflabs.emplace_back(newflab);
            }
            //保留扇叶状态
            if(flabellums.size() != 0) {
                for(int i=0;i<5;i++) {
                    candflabs[i].current_state = flabellums[i].current_state;
                }
            }
            flabellums = candflabs;
            flabellums[0].points2d = flabellum.points2d;
            // 保存四个角点  
            cv::Point2f points[4];
            flabellum.rotate.points(points);
            flabellum.points2d.resize(5);
            flabellum.points2d[0] = points[0];
            flabellum.points2d[1] = points[1];
            flabellum.points2d[2] = points[2];
            flabellum.points2d[3] = points[3];
            // 按左上左下右下右上的顺序
            std::sort(flabellum.points2d.begin(), flabellum.points2d.end()-1,[](cv::Point2f &a,cv::Point2f &b){return a.x < b.x;});
            std::sort(flabellum.points2d.begin(), flabellum.points2d.begin()+1,[](cv::Point2f &a,cv::Point2f &b){return a.y < b.y;});
            std::sort(flabellum.points2d.begin()+2, flabellum.points2d.end()-1,[](cv::Point2f &a,cv::Point2f &b){return a.y > b.y;});
            flabellum.points2d[4] = flabellum.rotate.center;
            targetarmor = flabellum;
        }

        void update(std::vector<Flabellum> flabs){
            //防止误识别 待数据稳定才更新
            trackstate = TrackState::LOST;
            if(flabs.size() != new_num) {
                update_time = 0;
                new_num = flabs.size();
            }
            else if(flabs.size() == new_num && update_time < 1) {
                update_time++;
            }
            else if(flabs.size() == new_num && update_time >= 1) {
                // iou匹配 更新扇叶的最新状态
                if(flabs.size() == 1) {
                    // 上一帧没识别到
                    flabs[0].current_state = 1;
                    flabellums.clear();
                    init(flabs.at(0));
                    lighted_num = 1;
                    trackstate = TrackState::TRACK;
                }
                else if(flabs.size() >= lighted_num && flabs.size() != 0) {
                    // 先和候选框匹配找出匹配的id
                    for(auto &flab: flabs) {
                        float max_iou = -1;
                        Flabellum fitted_flab;
                        for(auto &candflab: candflabs) {
                            float temp_iou = candflab ^ flab;
                            if(temp_iou > max_iou) {
                                fitted_flab = flab;
                                fitted_flab.id = candflab.id;
                                fitted_flab.current_state = 1;
                                fitted_flab.points2d = flab.points2d;
                                max_iou = temp_iou;
                            }
                        }
                        if(max_iou > -1) {
                            // 对对应id的扇叶更新状态
                            trackstate = TrackState::TRACK;
                            int chage_id = fitted_flab.id;
                            int old_state = flabellums[chage_id].current_state;
                            fitted_flab.points2d = flab.points2d;
                            flabellums[chage_id] = fitted_flab;
                            flabellums[chage_id].current_state = old_state;

                            //新增扇叶才需要更改状态
                            if(flabs.size() > lighted_num) {
                                flabellums[chage_id].current_state = old_state==2?2:old_state+1;
                                if(flabellums[chage_id].current_state == 1) {
                                    targetarmor = flabellums[chage_id];
                                }
                            }
                            
                            // 保持原来状态，有可能有刷新扇叶
                            if(flabs.size() == lighted_num) {
                                int target_id = 0;
                                // 查找原来是target的id
                                for(const auto &flabellum:flabellums) {
                                    if(flabellum.current_state == 1) {
                                        target_id = flabellum.id;
                                    }
                                }
                                if(old_state == 0 && chage_id != target_id) {
                                    // 扇叶刷新
                                    flabellums[chage_id].current_state = 1;
                                    flabellums[target_id].current_state = 0;
                                }
                                else {
                                    flabellums[chage_id].current_state = old_state;
                                }
                            }
                            //计算速度位置

                        }
                        else {
                            flab.current_state = 0;
                        }
                    }
                    if(flabs.size() > lighted_num) {
                        lighted_num++;
                    }
                }
            }
        }

        Flabellum target() {
            Flabellum temp;
            for(const auto &flab:flabellums) {
                if(flab.current_state == 1) {
                    return flab;
                }
            }
            return temp;
        }
    };

    // 定义圆的方程
    struct CircleModel {
        double operator()(const Eigen::Vector3d& p, const Eigen::VectorXd& x) const {
            return (p.head<3>() - x.head<3>()).norm() - 0.7;
        }
    };

    // 定义非线性最小二乘问题 求解圆心
    struct CircleFittingProblem {
        CircleFittingProblem(const std::vector<Eigen::Vector3d>& data_points)
            : data_points_(data_points) {}

        int operator()(const Eigen::VectorXd& x, Eigen::VectorXd& fvec) const {
            for (size_t i = 0; i < data_points_.size(); ++i) {
                fvec(i) = circle_model_(data_points_[i], x);
            }
            return 0;
        }

        int df(const Eigen::VectorXd& x, Eigen::MatrixXd& fjac) const {
            fjac.resize(data_points_.size(), 3);
            for (size_t i = 0; i < data_points_.size(); ++i) {
                const Eigen::Vector3d& p = data_points_[i];
                Eigen::Vector3d diff = p.head<3>() - x.head<3>();
                double dist = diff.norm();
                fjac(i, 0) = -diff[0] / dist; // derivative with respect to x coordinate of the circle center
                fjac(i, 1) = -diff[1] / dist; // derivative with respect to y coordinate of the circle center
                fjac(i, 2) = -diff[2] / dist; // derivative with respect to z coordinate of the circle center
                // fjac(i, 3) = -1.0; // derivative with respect to radius of the circle
            }
            return 0;
        }

        int inputs() const { return 3; } // 优化变量个数
        int values() const { return data_points_.size(); } // 残差个数

        const std::vector<Eigen::Vector3d>& data_points_;
        CircleModel circle_model_;
    };    

}  // namespace hnurm

#endif  // ARMOR_DETECTOR_ARMOR_HPP_
